from flask import Flask, jsonify, request
from flask_cors import CORS
from flask_socketio import SocketIO, emit
import numpy as np
import tensorflow as tf
import sqlite3

app = Flask(__name__)
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*")  # 添加在CORS之后

model = tf.keras.models.load_model('../ai_model/trained_models/iot_model.h5')

# 添加WebSocket路由
@socketio.on('connect')
def handle_connect():
    print('Client connected')

# 风险评分接口
@app.route('/api/risk', methods=['POST'])
def get_risk():
    try:
        data = request.get_json()
        
        # 输入校验（鲁棒性：防止非法输入）
        features = data.get('features', [])
        if not all(len(row) == 5 for row in features) or len(features) != 10:
            return jsonify({"code": 400, "error": "特征必须是 10 个长度为 5 的数组"})
        
        # 转换为 numpy 数组并调整形状
        features = np.array(features, dtype=np.float32).reshape(1, 10, 5)
        
        # 预测
        prediction = model.predict(features)
        risk_score = float(prediction[0][0]) * 100
        
        # 保留两位小数
        risk_score = round(risk_score, 2)
        
        return jsonify({"code": 200, "data": risk_score})
    
    except Exception as e:
        return jsonify({"code": 500, "error": f"服务器错误：{str(e)}"})
    
@app.route('/api/stats', methods=['GET'])
def get_stats():
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    c.execute("SELECT status, COUNT(*) FROM devices GROUP BY status")
    stats = {'online': 0, 'offline': 0}
    for row in c.fetchall():
        if row[0] == '在线':
            stats['online'] = row[1]
        else:
            stats['offline'] = row[1]

    c.execute("SELECT count FROM alerts WHERE id = 1")
    alerts = c.fetchone()[0]

    conn.close()
    return jsonify({
        "code": 200,
        "data": {
            "online": stats['online'],
            "offline": stats['offline'],
            "alerts": alerts,
            "risky": stats['offline']
        }
    })

@app.route('/api/devices')
def get_devices():
    status = request.args.getlist('status')
    risk_level = request.args.get('risk_level')
    start_time = request.args.get('start')
    end_time = request.args.get('end')
    
    query = "SELECT * FROM devices WHERE 1=1"
    params = []
    
    if status:
        query += " AND status IN ({})".format(','.join(['?']*len(status)))
        params.extend(status)
    
    # 示例风险等级过滤逻辑（需根据实际业务实现）
    if risk_level == 'high':
        query += " AND ip LIKE '192.168.1.1%'"
    
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    c.execute(query, params)
    devices = [{"id": row[0], "name": row[1], "ip": row[2], "status": row[3]} for row in c.fetchall()]
    conn.close()
    return jsonify({"code": 200, "data": devices})

@app.route('/api/event-types', methods=['GET'])
def get_event_types():
    return jsonify({
        "code": 200,
        "data": [
            {"value": "all", "label": "全部事件"},
            {"value": "ddos", "label": "DDoS攻击"},
            {"value": "login", "label": "异常登录"}
        ]
    })

@app.route('/api/risk-history/<event_type>')
def get_event_risk_history(event_type):
    # 模拟不同事件类型数据
    data_map = {
        "all": [45, 38, 52, 67, 41],
        "ddos": [30, 45, 28, 33, 40],
        "login": [15, 22, 18, 25, 20]
    }
    return jsonify({
        "code": 200,
        "data": data_map.get(event_type, [])
    })

# 新增威胁分析接口
@app.route('/api/threat-analysis', methods=['GET'])
def get_threat_analysis():
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    try:
        c.execute("SELECT type, count, trend FROM threat_analysis")
        rows = c.fetchall()
        types = [{"type": row[0], "count": row[1], "trend": row[2]} for row in rows]
        chartData = [row[1] for row in rows]  # 直接使用查询结果
        return jsonify({
            "code": 200,
            "data": {
                "types": types,
                "chartData": chartData
            }
        })
    except Exception as e:
        return jsonify({"code": 500, "error": str(e)})
    finally:
        conn.close()

@app.route('/api/security-events', methods=['GET'])
def get_security_events():
    limit = request.args.get('limit', default=10, type=int)
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    c.execute("SELECT time, description, level FROM security_events ORDER BY time DESC LIMIT ?", (limit,))
    events = [{
        "time": row[0],
        "description": row[1],
        "level": row[2]
    } for row in c.fetchall()]
    conn.close()
    return jsonify({"code": 200, "data": events})

@app.route('/api/notifications', methods=['GET'])
def get_notifications():
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    c.execute("SELECT title, time, type FROM notifications ORDER BY time DESC")
    notifications = [{
        "title": row[0],
        "time": row[1],
        "type": row[2]
    } for row in c.fetchall()]
    conn.close()
    return jsonify({"code": 200, "data": notifications})

@app.route('/api/risk-history', methods=['GET'])
def get_risk_history():
    try:
        # 模拟历史数据（实际应从数据库获取）
        data = [
            {"time": "2024-01-01", "score": 45},
            {"time": "2024-01-02", "score": 38},
            {"time": "2024-01-03", "score": 52},
            {"time": "2024-01-04", "score": 67},
            {"time": "2024-01-05", "score": 41}
        ]
        return jsonify({"code": 200, "data": data})
    except Exception as e:
        return jsonify({"code": 500, "error": str(e)})
    
@app.route('/api/devices/<int:device_id>')
def get_device(device_id):
    conn = sqlite3.connect('devices.db')
    c = conn.cursor()
    c.execute("SELECT * FROM devices WHERE id = ?", (device_id,))
    device = c.fetchone()
    conn.close()
    if device:
        return jsonify({
            "code": 200,
            "data": {
                "id": device[0],
                "name": device[1],
                "ip": device[2],
                "status": device[3]
            }
        })
    return jsonify({"code": 404, "error": "设备不存在"})

if __name__ == '__main__':
    socketio.run(app, host='0.0.0.0', port=5000, debug=True)